
import sys
sys.path.append('/Users/xbs/Code/HunterQuant')
import matplotlib.pyplot as plt
import mplfinance as mpf
import numpy as np
from PIL import Image
from data.data_module import DataModule
import pandas as pd
import matplotlib.dates as mdates
from pathlib import Path
from util.database import base_code_path
from util.stock_util import calc_positive_diff_dates


def plot_kline_to_image(ohlc_df, volume_df, save_path, dpi=300, figsize=(10,6)):
    
    # 配置样式
    mc = mpf.make_marketcolors(up='r', down='g', volume='in', 
                             edge='inherit', wick='inherit')
    s  = mpf.make_mpf_style(marketcolors=mc, gridstyle=':', y_on_right=True)
    
    # 绘图
    fig, axes = mpf.plot(ohlc_df, 
                        volume=True if volume_df is not None else False,
                        type='candle', 
                        style=s,
                        figsize=figsize,
                        returnfig=True,
                        closefig=False, # 修改为不自动关闭
                        )  
    for ax in axes:
        ax.set_xticklabels([])  # 清除Y轴数值标签
        ax.set_yticklabels([])  # 清除Y轴数值标签
        ax.yaxis.set_visible(False)  # 可选：彻底隐藏Y轴线
    # 正确获取主坐标轴
    #main_ax = axes[0]

    # 绘制颈线
    # neckline_points = [("2024-05-20", 6.0), ("2024-6-18", 6.0)]
    # line_style='--'
    # line_color='blue'
    # line_width=1.5
    # if neckline_points and len(neckline_points) == 2:
    #     # 转换日期格式
    #     dates = [mdates.date2num(pd.to_datetime(p[0])) for p in neckline_points]
    #     prices = [p[1] for p in neckline_points]
        
    #     # 在K线主图上绘制
    #     main_ax.plot(dates, prices, 
    #                 linestyle=line_style, 
    #                 color=line_color,
    #                 linewidth=line_width,
    #                 alpha=0.8,
    #                 label='Neckline')
        
    #     main_ax.set_xlim([ohlc_df.index.min(), ohlc_df.index.max()])
    #     # 添加图例
    #     main_ax.legend(loc='upper left')



    # 确保画布已渲染
    fig.set_dpi(dpi)
    fig.canvas.draw()

    # 获取实际画布尺寸
    w, h = fig.canvas.get_width_height()

    # 获取图像数据
    buf = fig.canvas.buffer_rgba()  # 替代tostring_argb()
    img = np.frombuffer(buf, dtype=np.uint8)
    img = img.reshape((h*2, w*2, 4))  # 使用实际尺寸
    
    # 转换为RGB并保存
    rgb_img = Image.fromarray(img[..., :3]).resize((3840,2160))
    rgb_img.save(save_path)
    plt.close(fig)

def mark_label(dm,row):
    code = row['code']
    pattern_end_date = row['pattern_end_date']
    end_date = calc_positive_diff_dates(code,is_index=False,date=pattern_end_date,delta_days=20)

    df_daily = dm.get_k_data(code, autype='qfq', begin_date=pattern_end_date, end_date=end_date)     
    if df_daily.empty:
        return False
    df_daily.set_index(['date'], inplace=True)

    max_value = df_daily['close'].max()
    max_index = df_daily['close'].idxmax()
    min_value = df_daily['close'].min()
    min_index = df_daily['close'].idxmin()
    # 买入、卖出日的开盘价
    beg_price = df_daily.loc[pattern_end_date]['close']

    raise_ratio = round(100*(max_value - beg_price)/beg_price,2)
    decrease_ratio = round(100*(min_value - beg_price)/beg_price,2)

    if(min_index < max_index and decrease_ratio >= 10):
        label = False
    else:
        if(raise_ratio >= 15 ):
            label =  True
        else:
            label =  False

    print(f"{code}:从{pattern_end_date}到{end_date}之间：在{max_index}出现最大涨幅{raise_ratio}%，在{min_index}出现最大跌幅{decrease_ratio}%,所以label:{label}")
    return label

def get_kline_data(dm,row):
    code = row['code']
    name = row['name']
    begin_date = row['pattern_begin_date']
    end_date = row['pattern_end_date']
    pattern_type = row['pattern_type']

    if(pattern_type == '双底'):
        dataset_path_prefix = "double_bottom"
    else:
         return
    
    df_daily = dm.get_k_data(code, autype='qfq', begin_date=begin_date, end_date=end_date)
    df_daily.rename(columns={'date':'Date','open': 'Open', 'close': 'Close','high': 'High','low': 'Low','volume': 'Volume'}, inplace=True)
    if df_daily.empty:
        return 
    
    label = mark_label(dm,row)
    if(label == False): 
        dataset_path_prefix = "other"      

    ohlc = df_daily
    ohlc['Date'] = pd.to_datetime(ohlc['Date'])
    ohlc = ohlc.set_index('Date')

    volume_df = df_daily['Volume']
    img_path = f"kline_dataset/train/{dataset_path_prefix}"
    img_name = f"{code}_{name}_{begin_date}_{end_date}"
    print(f"开始绘制:{img_name}")
    plot_kline_to_image(ohlc, volume_df, f'{img_path}/{img_name}.png')
    print(f"完成绘制:{img_name}")

    return 

def get_predict_kline_data(code,begin_date,end_date):
    dm = DataModule()   
    df_daily = dm.get_k_data(code, autype='qfq', begin_date=begin_date, end_date=end_date)
    df_daily.rename(columns={'date':'Date','open': 'Open', 'close': 'Close','high': 'High','low': 'Low','volume': 'Volume'}, inplace=True)

    return df_daily

def read_manual_pattern_info():
        file = Path(f"{base_code_path}/DL/PatternDataSet.xlsx")

        data_df = pd.read_excel(file, dtype=object)
        change_col = {
            "形态起始时间":'pattern_begin_date',
            "首次复盘发现突破颈线时间":'pattern_end_date',
            "代码":'code',
            "名称":'name',
            "形态底部值":'pattern_bottom_price',
            "形态顶部值":'pattern_top_price',
            "突破时的颈线位值":'pattern_neck_price',
            "趋势启动时间":'trending_begin_date',
            "趋势结束时间":'trending_end_date',
            "形态类型":'pattern_type',
            "二次追盘起始位置":'double_check_begin_date',
            "二次追盘结束位置":'double_check_end_date',
            "是否已标注":'isMark',
            "所属板块":'block',
            "换手率":'change_rate',
        }
        data_df.rename(columns=change_col,inplace=True)
        data_df['pattern_begin_date'] = data_df['pattern_begin_date'].astype(str)
        data_df['pattern_end_date'] = data_df['pattern_end_date'].astype(str)
        data_df['trending_begin_date'] = data_df['trending_begin_date'].astype(str)
        data_df['trending_end_date'] = data_df['trending_end_date'].astype(str)
        data_df['double_check_begin_date'] = data_df['double_check_begin_date'].astype(str)
        data_df['double_check_end_date'] = data_df['double_check_end_date'].astype(str)
        data_df['code'] = data_df['code'].astype(str)

        #print(data_df.head())

        return data_df

def gen_kline_dataset():
    data_df = read_manual_pattern_info()
    dm = DataModule()
    data_df.apply(lambda row: get_kline_data(dm,row), axis=1)


if __name__ == '__main__':
    gen_kline_dataset()
    #end_date = calc_positive_diff_dates("002241",is_index=False,date="2019-07-23",delta_days=20)
    #print(end_date)
